from unittest import TestCase

import pandas as pd
from sqlalchemy import create_engine


class PandasTest(TestCase):

    def test_pandas_sql_to_excel(self):
        engine = create_engine('postgresql://username:password@192.168.2.108:5432/demo')

        with engine.connect() as conn, conn.begin():
            data = pd.read_sql_table('user', conn)

        data.to_excel('a.xlsx')

        engine.dispose()

    def test_pandas_excel_to_dataframe(self):
        df = pd.read_excel('../resource/data.xlsx', index_col=0)

        # 增加新列
        df['sql'] = df.apply(
            lambda row: f"INSERT INTO price (name, age, price1) "
                        f"VALUES ('{row['name']}', {row['age']})",
            axis=1
        )
        print(df[['id', 'name', 'sql']])

    def test_1(self):
        s = pd.Series([22, 35, 58], name="年龄")

        print(s.dtypes)
        print(s.shape)

        print(s.info())
        print(s.describe())
        print(s.max())
        print(s.head())
        print(s.isin([35, 22]))
        print(s.notna())
        print(s.isna())
        print(s.isnull())

        print(s > 25)
        print(s > 25 & False)
        print(s > 25 & True)
        print(s > 25 | False)
        print(s > 25 | True)

    def test_2(self):
        df = pd.DataFrame([
            ['a', 35, '男'],
            ['b', 22, '未知'],
            ['c', 18, '女']
        ], columns=['名字', '年龄', '性别'])

        df = df.rename(columns={'年龄': 'age'})
        df = df.rename(columns={'年龄': 'age'}).rename(columns=str.upper)
        print(df)

        print(df['年龄'])
        print(df[['年龄', '性别']])
        print(df[df['年龄'] > 20])

        # 行, 列
        print(df.loc[df['年龄'] > 20, '年龄'])
        print(df.loc[df['年龄'] > 20, ['年龄', '性别']])

        print(df.iloc[0:3, 0:3])

        df.iloc[0:3, 2] = '狗'
        print(df.iloc[1])
        print(df.iloc[1, 2])

    def test_3(self):
        df = pd.DataFrame(
            {
                "Name": [
                    "Braund, Mr. Owen Harris",
                    "Allen, Mr. William Henry",
                    "Bonnell, Miss. Elizabeth",
                ],
                "Age": [22, 35, 58],
                "Sex": ["male", "male", "female"],
            }
        )

        # print(pd.Series([1, 2, 3], name='Age', dtype=np.str_))
        print(df.head(2))  # 查看前两行数据
        print(df.tail(2))  # 查看后两行数据
        print(df.dtypes)  # 查看每列的数据类型
        # df.to_excel('data.xlsx', index=False)

        print(df['Name'])
        print(df[['Name']])

        print(df['Age'] > 35)

        print(df[[False, False, True]])
